package com.nulldev.util.security.encryptAPI.impls.a51;

public class A51Cipher {

	private String key = null;
	static final int REG_X_LENGTH = 19;
	static final int REG_Y_LENGTH = 22;
	static final int REG_Z_LENGTH = 23;
	int[] regX = new int[REG_X_LENGTH];
	int[] regY = new int[REG_Y_LENGTH];
	int[] regZ = new int[REG_Z_LENGTH];

	public A51Cipher() {
	}

	public A51Cipher(final String key) {
		if (!this.setKey(key))
			throw new UnsupportedOperationException("Invalid key: " + key);
	}

	private void loadRegisters(String key) {
		for (int i = 0; i < REG_X_LENGTH; i++)
			regX[i] = Integer.parseInt(key.substring(i, i + 1));
		for (int i = 0; i < REG_Y_LENGTH; i++)
			regY[i] = Integer.parseInt(key.substring(REG_X_LENGTH + i, REG_X_LENGTH + i + 1));
		for (int i = 0; i < REG_Z_LENGTH; i++)
			regZ[i] = Integer.parseInt(key.substring(REG_X_LENGTH + REG_Y_LENGTH + i, REG_X_LENGTH + REG_Y_LENGTH + i + 1));
	}

	public String getKey() {
		return this.key;
	}

	public boolean setKey(String key) {
		if (key == null || key.isEmpty())
			return false;
		else if (key.length() == 64 && key.matches("[01]+")) {
			this.key = key;
			this.loadRegisters(key);
			return true;
		} else
			return false;
	}

	private int[] getKeystream(int length) {
		int[] regX = this.regX.clone();
		int[] regY = this.regY.clone();
		int[] regZ = this.regZ.clone();
		int[] keystream = new int[length];

		for (int i = 0; i < length; i++) {
			int maj = this.getMajority(regX[8], regY[10], regZ[10]);

			if (regX[8] == maj) {
				int newStart = regX[13] ^ regX[16] ^ regX[17] ^ regX[18];
				int[] temp = regX.clone();
				for (int j = 1; j < regX.length; j++)
					regX[j] = temp[j - 1];
				regX[0] = newStart;
			}

			if (regY[10] == maj) {
				int newStart = regY[20] ^ regY[21];
				int[] temp = regY.clone();
				for (int j = 1; j < regY.length; j++)
					regY[j] = temp[j - 1];
				regY[0] = newStart;
			}

			if (regZ[10] == maj) {
				int newStart = regZ[7] ^ regZ[20] ^ regZ[21] ^ regZ[22];
				int[] temp = regZ.clone();
				for (int j = 1; j < regZ.length; j++)
					regZ[j] = temp[j - 1];
				regZ[0] = newStart;
			}
			keystream[i] = regX[18] ^ regY[21] ^ regZ[22];
		}
		return keystream;
	}

	public String encrypt(String plaintext) {
		if (plaintext == null || plaintext.isEmpty())
			return null;
		StringBuilder s = new StringBuilder();
		int[] binary = this.toBinary(plaintext);
		int[] keystream = getKeystream(binary.length);
		for (int i = 0; i < binary.length; i++)
			s.append(binary[i] ^ keystream[i]);
		return s.toString();
	}

	public String decrypt(String ciphertext) {
		if (ciphertext == null || ciphertext.isEmpty())
			return null;
		StringBuilder s = new StringBuilder();
		int[] binary = new int[ciphertext.length()];
		int[] keystream = getKeystream(ciphertext.length());
		for (int i = 0; i < binary.length; i++) {
			binary[i] = Integer.parseInt(ciphertext.substring(i, i + 1));
			s.append(binary[i] ^ keystream[i]);
		}
		return this.toStr(s.toString());
	}

	private int[] toBinary(String text) {
		StringBuilder s = new StringBuilder();
		for (int i = 0; i < text.length(); i++) {
			String temp = Integer.toBinaryString(text.charAt(i));
			for (int j = temp.length(); j < 8; j++)
				temp = "0" + temp;
			s.append(temp);
		}
		String binaryStr = s.toString();
		int[] binary = new int[binaryStr.length()];
		for (int i = 0; i < binary.length; i++)
			binary[i] = Integer.parseInt(binaryStr.substring(i, i + 1));
		return binary;
	}

	private String toStr(String binary) {
		StringBuilder s = new StringBuilder();
		for (int i = 0; i <= binary.length() - 8; i += 8)
			s.append((char) Integer.parseInt(binary.substring(i, i + 8), 2));
		return s.toString();
	}

	private int getMajority(int x, int y, int z) {
		return x + y + z > 1 ? 1 : 0;
	}
}
